home *** CD-ROM | disk | FTP | other *** search
/ Personal Computer World 2008 February / PCWFEB08.iso / Software / Freeware / Miro 1.0 / Miro_Installer.exe / xulrunner / python / BitTorrent / StorageWrapper.py < prev    next >
Encoding:
Python Source  |  2007-11-12  |  16.9 KB  |  476 lines
  1. # Written by Bram Cohen
  2. # see LICENSE.txt for license information
  3.  
  4. from sha import sha
  5. from threading import Event
  6. from bitfield import Bitfield
  7.  
  8. def dummy_status(fractionDone = None, activity = None):
  9.     pass
  10.  
  11. def dummy_data_flunked(size):
  12.     pass
  13.  
  14. def dummy_hash_skip_func(block):
  15.     return False
  16.  
  17. class StorageWrapper:
  18.     def __init__(self, storage, request_size, hashes, 
  19.             piece_size, finished, failed, 
  20.             statusfunc = dummy_status, flag = Event(), check_hashes = True,
  21.             data_flunked = dummy_data_flunked,
  22.             hash_skip_func = dummy_hash_skip_func):
  23.         self.storage = storage
  24.         self.request_size = request_size
  25.         self.hashes = hashes
  26.         self.piece_size = piece_size
  27.         self.data_flunked = data_flunked
  28.         self.total_length = storage.get_total_length()
  29.         self.amount_left = self.total_length
  30.         if self.total_length <= piece_size * (len(hashes) - 1):
  31.             raise ValueError, 'bad data from tracker - total too small'
  32.         if self.total_length > piece_size * len(hashes):
  33.             raise ValueError, 'bad data from tracker - total too big'
  34.         self.finished = finished
  35.         self.failed = failed
  36.         self.numactive = [0] * len(hashes)
  37.         self.inactive_requests = [1] * len(hashes)
  38.         self.amount_inactive = self.total_length
  39.         self.endgame = False
  40.         self.have = Bitfield(len(hashes))
  41.         self.waschecked = [check_hashes] * len(hashes)
  42.         self.places = {}
  43.         self.holes = []
  44.         if len(hashes) == 0:
  45.             finished()
  46.             return
  47.         targets = {}
  48.         total = len(hashes)
  49.         for i in xrange(len(hashes)):
  50.             if not self._waspre(i):
  51.                 targets.setdefault(hashes[i], []).append(i)
  52.                 total -= 1
  53.         numchecked = 0.0
  54.         if total and check_hashes:
  55.             statusfunc({"activity" : 'checking existing file', 
  56.                 "fractionDone" : 0})
  57.         def markgot(piece, pos, self = self, check_hashes = check_hashes):
  58.             self.places[piece] = pos
  59.             self.have[piece] = True
  60.             self.amount_left -= self._piecelen(piece)
  61.             self.amount_inactive -= self._piecelen(piece)
  62.             self.inactive_requests[piece] = None
  63.             self.waschecked[piece] = check_hashes
  64.         lastlen = self._piecelen(len(hashes) - 1)
  65.         for i in xrange(len(hashes)):
  66.             files = self.storage.files_in_range(piece_size * i,
  67.                     self._piecelen(i))
  68.             if not self._waspre(i):
  69.                 self.holes.append(i)
  70.             elif not check_hashes or hash_skip_func(i, files):
  71.                 markgot(i, i)
  72.             else:
  73.                 sh = sha(self.storage.read(piece_size * i, lastlen))
  74.                 sp = sh.digest()
  75.                 sh.update(self.storage.read(piece_size * i + lastlen, self._piecelen(i) - lastlen))
  76.                 s = sh.digest()
  77.                 if s == hashes[i]:
  78.                     markgot(i, i)
  79.                 elif targets.get(s) and self._piecelen(i) == self._piecelen(targets[s][-1]):
  80.                     markgot(targets[s].pop(), i)
  81.                 elif not self.have[len(hashes) - 1] and sp == hashes[-1] and (i == len(hashes) - 1 or not self._waspre(len(hashes) - 1)):
  82.                     markgot(len(hashes) - 1, i)
  83.                 else:
  84.                     self.places[i] = i
  85.                 if flag.isSet():
  86.                     return
  87.                 numchecked += 1
  88.                 statusfunc({'fractionDone': 1 - float(self.amount_left) / self.total_length})
  89.         if self.amount_left == 0:
  90.             finished()
  91.  
  92.     def _waspre(self, piece):
  93.         return self.storage.was_preallocated(piece * self.piece_size, self._piecelen(piece))
  94.  
  95.     def _piecelen(self, piece):
  96.         if piece < len(self.hashes) - 1:
  97.             return self.piece_size
  98.         else:
  99.             return self.total_length - piece * self.piece_size
  100.  
  101.     def get_amount_left(self):
  102.         return self.amount_left
  103.  
  104.     def do_I_have_anything(self):
  105.         return self.amount_left < self.total_length
  106.  
  107.     def _make_inactive(self, index):
  108.         length = min(self.piece_size, self.total_length - self.piece_size * index)
  109.         l = []
  110.         x = 0
  111.         while x + self.request_size < length:
  112.             l.append((x, self.request_size))
  113.             x += self.request_size
  114.         l.append((x, length - x))
  115.         self.inactive_requests[index] = l
  116.  
  117.     def is_endgame(self):
  118.         return self.endgame
  119.  
  120.     def get_have_list(self):
  121.         return self.have.tostring()
  122.  
  123.     def do_I_have(self, index):
  124.         return self.have[index]
  125.  
  126.     def do_I_have_requests(self, index):
  127.         return not not self.inactive_requests[index]
  128.  
  129.     def new_request(self, index):
  130.         # returns (begin, length)
  131.         if self.inactive_requests[index] == 1:
  132.             self._make_inactive(index)
  133.         self.numactive[index] += 1
  134.         rs = self.inactive_requests[index]
  135.         r = min(rs)
  136.         rs.remove(r)
  137.         self.amount_inactive -= r[1]
  138.         if self.amount_inactive == 0:
  139.             self.endgame = True
  140.         return r
  141.  
  142.     def piece_came_in(self, index, begin, piece):
  143.         try:
  144.             return self._piece_came_in(index, begin, piece)
  145.         except IOError, e:
  146.             self.failed('IO Error ' + str(e))
  147.             return True
  148.  
  149.     def _piece_came_in(self, index, begin, piece):
  150.         if not self.places.has_key(index):
  151.             n = self.holes.pop(0)
  152.             if self.places.has_key(n):
  153.                 oldpos = self.places[n]
  154.                 old = self.storage.read(self.piece_size * oldpos, self._piecelen(n))
  155.                 if self.have[n] and sha(old).digest() != self.hashes[n]:
  156.                     self.failed('data corrupted on disk - maybe you have two copies running?')
  157.                     return True
  158.                 self.storage.write(self.piece_size * n, old)
  159.                 self.places[n] = n
  160.                 if index == oldpos or index in self.holes:
  161.                     self.places[index] = oldpos
  162.                 else:
  163.                     for p, v in self.places.items():
  164.                         if v == index:
  165.                             break
  166.                     self.places[index] = index
  167.                     self.places[p] = oldpos
  168.                     old = self.storage.read(self.piece_size * index, self.piece_size)
  169.                     self.storage.write(self.piece_size * oldpos, old)
  170.             elif index in self.holes or index == n:
  171.                 if not self._waspre(n):
  172.                     self.storage.write(self.piece_size * n, self._piecelen(n) * chr(0xFF))
  173.                 self.places[index] = n
  174.             else:
  175.                 for p, v in self.places.items():
  176.                     if v == index:
  177.                         break
  178.                 self.places[index] = index
  179.                 self.places[p] = n
  180.                 old = self.storage.read(self.piece_size * index, self._piecelen(n))
  181.                 self.storage.write(self.piece_size * n, old)
  182.         self.storage.write(self.places[index] * self.piece_size + begin, piece)
  183.         self.numactive[index] -= 1
  184.         if not self.inactive_requests[index] and not self.numactive[index]:
  185.             if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() == self.hashes[index]:
  186.                 self.have[index] = True
  187.                 self.inactive_requests[index] = None
  188.                 self.waschecked[index] = True
  189.                 self.amount_left -= self._piecelen(index)
  190.                 if self.amount_left == 0:
  191.                     self.finished()
  192.             else:
  193.                 self.data_flunked(self._piecelen(index))
  194.                 self.inactive_requests[index] = 1
  195.                 self.amount_inactive += self._piecelen(index)
  196.                 return False
  197.         return True
  198.  
  199.     def request_lost(self, index, begin, length):
  200.         self.inactive_requests[index].append((begin, length))
  201.         self.amount_inactive += length
  202.         self.numactive[index] -= 1
  203.  
  204.     def get_piece(self, index, begin, length):
  205.         try:
  206.             return self._get_piece(index, begin, length)
  207.         except IOError, e:
  208.             self.failed('IO Error ' + str(e))
  209.             return None
  210.  
  211.     def _get_piece(self, index, begin, length):
  212.         if not self.have[index]:
  213.             return None
  214.         if not self.waschecked[index]:
  215.             if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() != self.hashes[index]:
  216.                 self.failed('told file complete on start-up, but piece failed hash check')
  217.                 return None
  218.             self.waschecked[index] = True
  219.         if begin + length > self._piecelen(index):
  220.             return None
  221.         return self.storage.read(self.piece_size * self.places[index] + begin, length)
  222.  
  223. class DummyStorage:
  224.     def __init__(self, total, pre = False, ranges = []):
  225.         self.pre = pre
  226.         self.ranges = ranges
  227.         self.s = chr(0xFF) * total
  228.         self.done = False
  229.  
  230.     def was_preexisting(self):
  231.         return self.pre
  232.  
  233.     def was_preallocated(self, begin, length):
  234.         for b, l in self.ranges:
  235.             if begin >= b and begin + length <= b + l:
  236.                 return True
  237.         return False
  238.  
  239.     def get_total_length(self):
  240.         return len(self.s)
  241.  
  242.     def read(self, begin, length):
  243.         return self.s[begin:begin + length]
  244.  
  245.     def write(self, begin, piece):
  246.         self.s = self.s[:begin] + piece + self.s[begin + len(piece):]
  247.  
  248.     def finished(self):
  249.         self.done = True
  250.  
  251. def test_basic():
  252.     ds = DummyStorage(3)
  253.     sw = StorageWrapper(ds, 2, [sha('abc').digest()], 4, ds.finished, None)
  254.     assert sw.get_amount_left() == 3
  255.     assert not sw.do_I_have_anything()
  256.     assert sw.get_have_list() == chr(0)
  257.     assert sw.do_I_have_requests(0)
  258.     x = []
  259.     x.append(sw.new_request(0))
  260.     assert sw.do_I_have_requests(0)
  261.     x.append(sw.new_request(0))
  262.     assert not sw.do_I_have_requests(0)
  263.     x.sort()
  264.     assert x == [(0, 2), (2, 1)]
  265.     sw.request_lost(0, 2, 1)
  266.     del x[-1]
  267.     assert sw.do_I_have_requests(0)
  268.     x.append(sw.new_request(0))
  269.     assert x == [(0, 2), (2, 1)]
  270.     assert not sw.do_I_have_requests(0)
  271.     sw.piece_came_in(0, 0, 'ab')
  272.     assert not sw.do_I_have_requests(0)
  273.     assert sw.get_amount_left() == 3
  274.     assert not sw.do_I_have_anything()
  275.     assert sw.get_have_list() == chr(0)
  276.     assert not ds.done
  277.     sw.piece_came_in(0, 2, 'c')
  278.     assert not sw.do_I_have_requests(0)
  279.     assert sw.get_amount_left() == 0
  280.     assert sw.do_I_have_anything()
  281.     assert sw.get_have_list() == chr(0x80)
  282.     assert sw.get_piece(0, 0, 3) == 'abc'
  283.     assert sw.get_piece(0, 1, 2) == 'bc'
  284.     assert sw.get_piece(0, 0, 2) == 'ab'
  285.     assert sw.get_piece(0, 1, 1) == 'b'
  286.     assert ds.done
  287.  
  288. def test_two_pieces():
  289.     ds = DummyStorage(4)
  290.     sw = StorageWrapper(ds, 3, [sha('abc').digest(),
  291.         sha('d').digest()], 3, ds.finished, None)
  292.     assert sw.get_amount_left() == 4
  293.     assert not sw.do_I_have_anything()
  294.     assert sw.get_have_list() == chr(0)
  295.     assert sw.do_I_have_requests(0)
  296.     assert sw.do_I_have_requests(1)
  297.  
  298.     assert sw.new_request(0) == (0, 3)
  299.     assert sw.get_amount_left() == 4
  300.     assert not sw.do_I_have_anything()
  301.     assert sw.get_have_list() == chr(0)
  302.     assert not sw.do_I_have_requests(0)
  303.     assert sw.do_I_have_requests(1)
  304.  
  305.     assert sw.new_request(1) == (0, 1)
  306.     assert sw.get_amount_left() == 4
  307.     assert not sw.do_I_have_anything()
  308.     assert sw.get_have_list() == chr(0)
  309.     assert not sw.do_I_have_requests(0)
  310.     assert not sw.do_I_have_requests(1)
  311.  
  312.     sw.piece_came_in(0, 0, 'abc')
  313.     assert sw.get_amount_left() == 1
  314.     assert sw.do_I_have_anything()
  315.     assert sw.get_have_list() == chr(0x80)
  316.     assert not sw.do_I_have_requests(0)
  317.     assert not sw.do_I_have_requests(1)
  318.     assert sw.get_piece(0, 0, 3) == 'abc'
  319.     assert not ds.done
  320.  
  321.     sw.piece_came_in(1, 0, 'd')
  322.     assert ds.done
  323.     assert sw.get_amount_left() == 0
  324.     assert sw.do_I_have_anything()
  325.     assert sw.get_have_list() == chr(0xC0)
  326.     assert not sw.do_I_have_requests(0)
  327.     assert not sw.do_I_have_requests(1)
  328.     assert sw.get_piece(1, 0, 1) == 'd'
  329.  
  330. def test_hash_fail():
  331.     ds = DummyStorage(4)
  332.     sw = StorageWrapper(ds, 4, [sha('abcd').digest()], 4, ds.finished, None)
  333.     assert sw.get_amount_left() == 4
  334.     assert not sw.do_I_have_anything()
  335.     assert sw.get_have_list() == chr(0)
  336.     assert sw.do_I_have_requests(0)
  337.  
  338.     assert sw.new_request(0) == (0, 4)
  339.     sw.piece_came_in(0, 0, 'abcx')
  340.     assert sw.get_amount_left() == 4
  341.     assert not sw.do_I_have_anything()
  342.     assert sw.get_have_list() == chr(0)
  343.     assert sw.do_I_have_requests(0)
  344.  
  345.     assert sw.new_request(0) == (0, 4)
  346.     assert not ds.done
  347.     sw.piece_came_in(0, 0, 'abcd')
  348.     assert ds.done
  349.     assert sw.get_amount_left() == 0
  350.     assert sw.do_I_have_anything()
  351.     assert sw.get_have_list() == chr(0x80)
  352.     assert not sw.do_I_have_requests(0)
  353.  
  354. def test_lazy_hashing():
  355.     ds = DummyStorage(4, ranges = [(0, 4)])
  356.     flag = Event()
  357.     sw = StorageWrapper(ds, 4, [sha('abcd').digest()], 4, ds.finished, lambda x, flag = flag: flag.set(), check_hashes = False)
  358.     assert sw.get_piece(0, 0, 2) is None
  359.     assert flag.isSet()
  360.  
  361. def test_lazy_hashing_pass():
  362.     ds = DummyStorage(4)
  363.     flag = Event()
  364.     sw = StorageWrapper(ds, 4, [sha(chr(0xFF) * 4).digest()], 4, ds.finished, lambda x, flag = flag: flag.set(), check_hashes = False)
  365.     assert sw.get_piece(0, 0, 2) is None
  366.     assert not flag.isSet()
  367.  
  368. def test_preexisting():
  369.     ds = DummyStorage(4, True, [(0, 4)])
  370.     sw = StorageWrapper(ds, 2, [sha(chr(0xFF) * 2).digest(), 
  371.         sha('ab').digest()], 2, ds.finished, None)
  372.     assert sw.get_amount_left() == 2
  373.     assert sw.do_I_have_anything()
  374.     assert sw.get_have_list() == chr(0x80)
  375.     assert not sw.do_I_have_requests(0)
  376.     assert sw.do_I_have_requests(1)
  377.     assert sw.new_request(1) == (0, 2)
  378.     assert not ds.done
  379.     sw.piece_came_in(1, 0, 'ab')
  380.     assert ds.done
  381.     assert sw.get_amount_left() == 0
  382.     assert sw.do_I_have_anything()
  383.     assert sw.get_have_list() == chr(0xC0)
  384.     assert not sw.do_I_have_requests(0)
  385.     assert not sw.do_I_have_requests(1)
  386.  
  387. def test_total_too_short():
  388.     ds = DummyStorage(4)
  389.     try:
  390.         StorageWrapper(ds, 4, [sha(chr(0xff) * 4).digest(),
  391.             sha(chr(0xFF) * 4).digest()], 4, ds.finished, None)
  392.         raise 'fail'
  393.     except ValueError:
  394.         pass
  395.  
  396. def test_total_too_big():
  397.     ds = DummyStorage(9)
  398.     try:
  399.         sw = StorageWrapper(ds, 4, [sha('qqqq').digest(),
  400.             sha(chr(0xFF) * 4).digest()], 4, ds.finished, None)
  401.         raise 'fail'
  402.     except ValueError:
  403.         pass
  404.  
  405. def test_end_above_total_length():
  406.     ds = DummyStorage(3, True)
  407.     sw = StorageWrapper(ds, 4, [sha('qqq').digest()], 4, ds.finished, None)
  408.     assert sw.get_piece(0, 0, 4) == None
  409.  
  410. def test_end_past_piece_end():
  411.     ds = DummyStorage(4, True, ranges = [(0, 4)])
  412.     sw = StorageWrapper(ds, 4, [sha(chr(0xFF) * 2).digest(), 
  413.         sha(chr(0xFF) * 2).digest()], 2, ds.finished, None)
  414.     assert ds.done
  415.     assert sw.get_piece(0, 0, 3) == None
  416.  
  417. from random import shuffle
  418.  
  419. def test_alloc_random():
  420.     ds = DummyStorage(101)
  421.     sw = StorageWrapper(ds, 1, [sha(chr(i)).digest() for i in xrange(101)], 1, ds.finished, None)
  422.     for i in xrange(100):
  423.         assert sw.new_request(i) == (0, 1)
  424.     r = range(100)
  425.     shuffle(r)
  426.     for i in r:
  427.         sw.piece_came_in(i, 0, chr(i))
  428.     for i in xrange(100):
  429.         assert sw.get_piece(i, 0, 1) == chr(i)
  430.     assert ds.s[:100] == ''.join([chr(i) for i in xrange(100)])
  431.  
  432. def test_alloc_resume():
  433.     ds = DummyStorage(101)
  434.     sw = StorageWrapper(ds, 1, [sha(chr(i)).digest() for i in xrange(101)], 1, ds.finished, None)
  435.     for i in xrange(100):
  436.         assert sw.new_request(i) == (0, 1)
  437.     r = range(100)
  438.     shuffle(r)
  439.     for i in r[:50]:
  440.         sw.piece_came_in(i, 0, chr(i))
  441.     assert ds.s[50:] == chr(0xFF) * 51
  442.     ds.ranges = [(0, 50)]
  443.     sw = StorageWrapper(ds, 1, [sha(chr(i)).digest() for i in xrange(101)], 1, ds.finished, None)
  444.     for i in r[50:]:
  445.         sw.piece_came_in(i, 0, chr(i))
  446.     assert ds.s[:100] == ''.join([chr(i) for i in xrange(100)])
  447.  
  448. def test_last_piece_pre():
  449.     ds = DummyStorage(3, ranges = [(2, 1)])
  450.     ds.s = chr(0xFF) + chr(0xFF) + 'c'
  451.     sw = StorageWrapper(ds, 2, [sha('ab').digest(), sha('c').digest()], 2, ds.finished, None)
  452.     assert not sw.do_I_have_requests(1)
  453.     assert sw.do_I_have_requests(0)
  454.  
  455. def test_not_last_pre():
  456.     ds = DummyStorage(3, ranges = [(1, 1)])
  457.     ds.s = chr(0xFF) + 'a' + chr(0xFF)
  458.     sw = StorageWrapper(ds, 1, [sha('a').digest()] * 3, 1, ds.finished, None)
  459.     assert not sw.do_I_have_requests(1)
  460.     assert sw.do_I_have_requests(0)
  461.     assert sw.do_I_have_requests(2)
  462.  
  463. def test_last_piece_not_pre():
  464.     ds = DummyStorage(51, ranges = [(50, 1)])
  465.     sw = StorageWrapper(ds, 2, [sha('aa').digest()] * 25 + [sha('b').digest()], 2, ds.finished, None)
  466.     for i in xrange(25):
  467.         assert sw.new_request(i) == (0, 2)
  468.     assert sw.new_request(25) == (0, 1)
  469.     sw.piece_came_in(25, 0, 'b')
  470.     r = range(25)
  471.     shuffle(r)
  472.     for i in r:
  473.         sw.piece_came_in(i, 0, 'aa')
  474.     assert ds.done
  475.     assert ds.s == 'a' * 50 + 'b'
  476.